#!/bin/bash

# Set default values
MODEL_NAME="bigcode/starcoderbase-1b"  # unsloth/Meta-Llama-3.1-8B
LANG="rkt"
ROOT_DATASET="humaneval"
TEMPERATURE="0.2"
NUM_GPUS="1"
TEMPERATURE="0.2"
BATCH_SIZE="20"
COMPLETION_LIMIT="20"
OUTPUT_DIR_PREFIX="multiple_results"
NUM_GPUS=2
TAG=""  # New default value for tag

# RAG config
KNOWLEDGE_PATH="documents/knowledge/racket/racket_knowledge.jsonl"
KNOWLEDGE_INDEX_CACHE_PATH="rag/index/"
KNOWLEDGE_MODEL_NAME="jinaai/jina-embeddings-v2-base-code"

CODE_DATA_PATH="/path/to/code/dataset"
CODE_INDEX_CACHE_PATH="/path/to/code/index/cache"
CODE_MODEL_NAME="code-embedding-model-name"


# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --name)
            MODEL_NAME="$2"
            shift 2
            ;;
        --revision)
            REVISION="$2"
            shift 2
            ;;
        --lang)
            LANG="$2"
            shift 2
            ;;
        --root-dataset)
            ROOT_DATASET="$2"
            shift 2
            ;;
        --temperature)
            TEMPERATURE="$2"
            shift 2
            ;;
        --num-gpus)
            NUM_GPUS="$2"
            shift 2
            ;;
        --tag)
            TAG="$2"
            shift 2
            ;;
        --knowledge-path)
            KNOWLEDGE_PATH="$2"
            shift 2
            ;;
        --knowledge-index-cache-path)
            KNOWLEDGE_INDEX_CACHE_PATH="$2"
            shift 2
            ;;
        --knowledge-model-name)
            KNOWLEDGE_MODEL_NAME="$2"
            shift 2
            ;;
        --code-data-path)
            CODE_DATA_PATH="$2"
            shift 2
            ;;
        --code-index-cache-path)
            CODE_INDEX_CACHE_PATH="$2"
            shift 2
            ;;
        --code-model-name)
            CODE_MODEL_NAME="$2"
            shift 2
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

# Check if required arguments are provided
if [ -z "$LANG" ] || [ -z "$ROOT_DATASET" ] || [ -z "$TEMPERATURE" ]; then
    echo "Error: --lang, --root-dataset, and --temperature are required arguments."
    exit 1
fi

# Set output directory
OUTPUT_DIR="$OUTPUT_DIR_PREFIX/$ROOT_DATASET/${MODEL_NAME//\//_}/$LANG/$TAG"
KNOWLEDGE_INDEX_CACHE_PATH="$KNOWLEDGE_INDEX_CACHE_PATH/knowledge_$LANG.index"

# Create the output directory if it doesn't exist
mkdir -p "$OUTPUT_DIR"

# Run the Python script
# Check if the output directory already contains results
# if [ -n "$(ls -A "$OUTPUT_DIR")" ]; then
#     echo "Output directory $OUTPUT_DIR already contains results. Skipping inference step."
# else

echo "Running inference..."
python inference_rag.py \
    --name "$MODEL_NAME" \
    --lang "$LANG" \
    --root-dataset "$ROOT_DATASET" \
    --temperature "$TEMPERATURE" \
    --output-dir "$OUTPUT_DIR" \
    --batch-size "$BATCH_SIZE" \
    --completion-limit "$COMPLETION_LIMIT" \
    --num_gpus "$NUM_GPUS" \
    --knowledge-path "$KNOWLEDGE_PATH" \
    --knowledge-index-cache-path "$KNOWLEDGE_INDEX_CACHE_PATH" \
    --knowledge-model-name "$KNOWLEDGE_MODEL_NAME" \
    # --code-data-path "$CODE_DATA_PATH" \
    # --code-index-cache-path "$CODE_INDEX_CACHE_PATH" \
    # --code-model-name "$CODE_MODEL_NAME"

# # run
# python multipl_e/evaluation/src/main.py --dir "$OUTPUT_DIR" --output-dir "$OUTPUT_DIR"

# # metrics
# python multipl_e/pass_k.py "$OUTPUT_DIR"